from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import json

import io

def read_data(filename):
    with io.open(filename, encoding="utf-8") as f:
        data = f.read()
    data = list(data)
    return data



def build_dataset(words, n_words):
    
    count = [['UNK', -1]]                                
    count.extend(collections.Counter(words).most_common(n_words - 1))  
                                                                       
    dictionary = dict()                                                
    for word, _ in count:
        dictionary[word] = len(dictionary)                             

    data = list()
    unk_count = 0
    for word in words:                                                
        index = dictionary.get(word, 0)                                
        if index == 0:   
            unk_count += 1
        data.append(index)
    count[0][1] = unk_count

    reversed_dictionary = dict(zip(dictionary.values(), dictionary.keys()))  
    
    return data, count, dictionary, reversed_dictionary  


filename = './QuanSongCi.txt'

vocabulary = read_data(filename)
print('Data size', len(vocabulary))

vocabulary_size = 5000

data, count, dictionary, reversed_dictionary = build_dataset(vocabulary,vocabulary_size - 1)


with io.open("./dictionary.json","w",encoding='utf-8') as f:
	f.write(unicode(json.dumps(dictionary, ensure_ascii=False)))

with io.open("./reverse_dictionary.json","w",encoding='utf-8') as f:
	f.write(unicode(json.dumps(reversed_dictionary, ensure_ascii=False)))
    

